"""
WAL (Write-Ahead Logging) 管理器
"""

import os
import sys
import json
import threading
import time
from typing import List, Dict, Any, Optional, BinaryIO
from datetime import datetime
import struct

# 添加当前目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

from wal_types import (
    WALRecord, WALRecordType, WALControlInfo, CheckpointInfo, WALConstants
)
from statistics import Logger


class WALManager:
    """WAL管理器 - 负责WAL的写入、管理和恢复"""
    
    def __init__(self, database_path: str, wal_file_size: int = WALConstants.DEFAULT_WAL_FILE_SIZE):
        """
        初始化WAL管理器
        
        Args:
            database_path: 数据库目录路径
            wal_file_size: WAL文件大小限制
        """
        self.database_path = database_path
        self.wal_file_size = wal_file_size
        self.logger = Logger()
        
        # WAL目录
        self.wal_directory = os.path.join(database_path, WALConstants.WAL_DIRECTORY)
        os.makedirs(self.wal_directory, exist_ok=True)
        
        # WAL文件路径
        self.control_file_path = os.path.join(self.wal_directory, WALConstants.WAL_CONTROL_FILE)
        self.checkpoint_file_path = os.path.join(self.wal_directory, WALConstants.WAL_CHECKPOINT_FILE)
        
        # WAL状态
        self.current_lsn = 1
        self.current_wal_file_number = 1
        self.current_wal_file: Optional[BinaryIO] = None
        self.current_wal_file_size = 0
        
        # 线程安全
        self.wal_lock = threading.RLock()
        self.flush_lock = threading.Lock()
        
        # WAL缓冲区
        self.wal_buffer: List[WALRecord] = []
        self.buffer_size_limit = 100  # 缓冲区记录数限制
        
        # 活跃事务
        self.active_transactions: Dict[str, List[int]] = {}  # transaction_id -> [lsn_list]
        
        # 初始化
        self._load_control_info()
        self._open_current_wal_file()
        
        self.logger.info("WAL_INIT", 0, f"WAL Manager initialized at {self.wal_directory}")
    
    def _load_control_info(self):
        """加载WAL控制信息"""
        try:
            if os.path.exists(self.control_file_path):
                with open(self.control_file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                self.current_lsn = data.get('next_lsn', 1)
                self.current_wal_file_number = int(data.get('current_wal_file', 'wal_000001.log')[4:10])
                
                self.logger.info("WAL_LOAD", 0, f"Loaded WAL control info: LSN={self.current_lsn}")
            else:
                # 创建初始控制文件
                self._save_control_info()
                self.logger.info("WAL_INIT", 0, "Created initial WAL control file")
                
        except Exception as e:
            self.logger.error("WAL_LOAD", 0, f"Failed to load WAL control info: {e}")
            # 使用默认值
            self.current_lsn = 1
            self.current_wal_file_number = 1
    
    def _save_control_info(self):
        """保存WAL控制信息"""
        try:
            control_info = {
                'current_wal_file': f"wal_{self.current_wal_file_number:06d}.log",
                'next_lsn': self.current_lsn,
                'last_checkpoint_lsn': self._get_last_checkpoint_lsn(),
                'last_checkpoint_time': self._get_last_checkpoint_time(),
                'wal_file_size': self.wal_file_size,
                'database_version': "1.0.0"
            }
            
            with open(self.control_file_path, 'w', encoding='utf-8') as f:
                json.dump(control_info, f, indent=2, default=str)
                
        except Exception as e:
            self.logger.error("WAL_SAVE", 0, f"Failed to save WAL control info: {e}")
    
    def _get_current_wal_file_path(self) -> str:
        """获取当前WAL文件路径"""
        filename = f"wal_{self.current_wal_file_number:06d}.log"
        return os.path.join(self.wal_directory, filename)
    
    def _open_current_wal_file(self):
        """打开当前WAL文件用于写入"""
        try:
            if self.current_wal_file:
                self.current_wal_file.close()
            
            wal_file_path = self._get_current_wal_file_path()
            
            # 如果文件存在，获取当前大小
            if os.path.exists(wal_file_path):
                self.current_wal_file_size = os.path.getsize(wal_file_path)
            else:
                self.current_wal_file_size = 0
            
            # 以追加模式打开
            self.current_wal_file = open(wal_file_path, 'ab')
            
            self.logger.info("WAL_OPEN", 0, f"Opened WAL file: {wal_file_path}")
            
        except Exception as e:
            self.logger.error("WAL_OPEN", 0, f"Failed to open WAL file: {e}")
            raise
    
    def _rotate_wal_file(self):
        """轮转WAL文件"""
        try:
            if self.current_wal_file:
                self.current_wal_file.close()
            
            # 切换到下一个WAL文件
            self.current_wal_file_number += 1
            self.current_wal_file_size = 0
            
            self._open_current_wal_file()
            self._save_control_info()
            
            self.logger.info("WAL_ROTATE", 0, f"Rotated to WAL file: {self.current_wal_file_number}")
            
        except Exception as e:
            self.logger.error("WAL_ROTATE", 0, f"Failed to rotate WAL file: {e}")
            raise
    
    def write_transaction_begin(self, transaction_id: str) -> int:
        """写入事务开始记录"""
        return self._write_wal_record(
            transaction_id=transaction_id,
            record_type=WALRecordType.BEGIN,
            table_name=None,
            page_id=None,
            before_image=None,
            after_image=None
        )
    
    def write_transaction_commit(self, transaction_id: str) -> int:
        """写入事务提交记录"""
        lsn = self._write_wal_record(
            transaction_id=transaction_id,
            record_type=WALRecordType.COMMIT,
            table_name=None,
            page_id=None,
            before_image=None,
            after_image=None
        )
        
        # 从活跃事务中移除
        with self.wal_lock:
            if transaction_id in self.active_transactions:
                del self.active_transactions[transaction_id]
        
        return lsn
    
    def write_transaction_rollback(self, transaction_id: str) -> int:
        """写入事务回滚记录"""
        lsn = self._write_wal_record(
            transaction_id=transaction_id,
            record_type=WALRecordType.ROLLBACK,
            table_name=None,
            page_id=None,
            before_image=None,
            after_image=None
        )
        
        # 从活跃事务中移除
        with self.wal_lock:
            if transaction_id in self.active_transactions:
                del self.active_transactions[transaction_id]
        
        return lsn
    
    def write_data_change(self, transaction_id: str, operation: str, table_name: str,
                         page_id: int = None, before_image: bytes = None, 
                         after_image: bytes = None, offset: int = None) -> int:
        """
        写入数据变更记录
        
        Args:
            transaction_id: 事务ID
            operation: 操作类型 (INSERT/UPDATE/DELETE)
            table_name: 表名
            page_id: 页面ID
            before_image: 修改前数据
            after_image: 修改后数据
            offset: 页内偏移
        """
        record_type_map = {
            'INSERT': WALRecordType.INSERT,
            'UPDATE': WALRecordType.UPDATE,
            'DELETE': WALRecordType.DELETE
        }
        
        record_type = record_type_map.get(operation.upper())
        if not record_type:
            raise ValueError(f"Unsupported operation: {operation}")
        
        return self._write_wal_record(
            transaction_id=transaction_id,
            record_type=record_type,
            table_name=table_name,
            page_id=page_id,
            before_image=before_image,
            after_image=after_image,
            offset=offset
        )
    
    def _write_wal_record(self, transaction_id: str, record_type: WALRecordType,
                         table_name: Optional[str] = None, page_id: Optional[int] = None,
                         before_image: Optional[bytes] = None, after_image: Optional[bytes] = None,
                         offset: Optional[int] = None) -> int:
        """写入WAL记录"""
        with self.wal_lock:
            # 分配LSN
            lsn = self.current_lsn
            self.current_lsn += 1
            
            # 创建WAL记录
            wal_record = WALRecord(
                lsn=lsn,
                transaction_id=transaction_id,
                record_type=record_type,
                timestamp=datetime.now(),
                table_name=table_name,
                page_id=page_id,
                offset=offset,
                before_image=before_image,
                after_image=after_image
            )
            
            # 添加到缓冲区
            self.wal_buffer.append(wal_record)
            
            # 跟踪活跃事务
            if record_type in [WALRecordType.BEGIN, WALRecordType.INSERT, 
                              WALRecordType.UPDATE, WALRecordType.DELETE]:
                if transaction_id not in self.active_transactions:
                    self.active_transactions[transaction_id] = []
                self.active_transactions[transaction_id].append(lsn)
            
            # 检查是否需要刷新缓冲区
            if len(self.wal_buffer) >= self.buffer_size_limit:
                self._flush_buffer()
            
            self.logger.info("WAL_WRITE", lsn, f"Written WAL record: {record_type.name} for tx {transaction_id}")
            
            return lsn
    
    def _flush_buffer(self):
        """刷新WAL缓冲区到磁盘"""
        if not self.wal_buffer:
            return
        
        try:
            for record in self.wal_buffer:
                serialized_data = record.serialize()
                
                # 检查是否需要轮转文件
                if self.current_wal_file_size + len(serialized_data) > self.wal_file_size:
                    self._rotate_wal_file()
                
                # 写入数据
                self.current_wal_file.write(serialized_data)
                self.current_wal_file_size += len(serialized_data)
            
            # 强制刷新到磁盘
            self.current_wal_file.flush()
            os.fsync(self.current_wal_file.fileno())
            
            # 清空缓冲区
            flushed_count = len(self.wal_buffer)
            self.wal_buffer.clear()
            
            self.logger.info("WAL_FLUSH", 0, f"Flushed {flushed_count} WAL records to disk")
            
        except Exception as e:
            self.logger.error("WAL_FLUSH", 0, f"Failed to flush WAL buffer: {e}")
            raise
    
    def force_flush(self):
        """强制刷新WAL到磁盘"""
        with self.flush_lock:
            self._flush_buffer()
            self._save_control_info()
    
    def get_current_lsn(self) -> int:
        """获取当前LSN"""
        return self.current_lsn
    
    def get_flush_lsn(self) -> int:
        """获取已刷新的LSN（用于确定哪些页面可以安全刷新）"""
        # 简化实现：返回当前LSN-1，实际应该跟踪已刷新的LSN
        return max(1, self.current_lsn - len(self.wal_buffer) - 1)
    
    def list_wal_files(self) -> List[str]:
        """列出所有WAL文件"""
        try:
            files = []
            for filename in os.listdir(self.wal_directory):
                if filename.startswith(WALConstants.WAL_FILE_PREFIX) and filename.endswith(WALConstants.WAL_FILE_EXTENSION):
                    files.append(filename)
            return sorted(files)
        except Exception as e:
            self.logger.error("WAL_LIST", 0, f"Failed to list WAL files: {e}")
            return []
    
    def get_total_wal_size(self) -> int:
        """获取WAL文件总大小"""
        total_size = 0
        try:
            for filename in self.list_wal_files():
                file_path = os.path.join(self.wal_directory, filename)
                if os.path.exists(file_path):
                    total_size += os.path.getsize(file_path)
        except Exception as e:
            self.logger.error("WAL_SIZE", 0, f"Failed to calculate WAL size: {e}")
        return total_size
    
    def get_pending_write_count(self) -> int:
        """获取待写入的记录数"""
        return len(self.wal_buffer)
    
    def get_write_rate(self) -> float:
        """获取写入速率（记录/秒）"""
        # 简化实现，实际应该维护统计信息
        return 0.0
    
    def _get_last_checkpoint_lsn(self) -> int:
        """获取最后检查点LSN"""
        try:
            if os.path.exists(self.checkpoint_file_path):
                with open(self.checkpoint_file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                return data.get('checkpoint_lsn', 0)
        except:
            pass
        return 0
    
    def _get_last_checkpoint_time(self) -> str:
        """获取最后检查点时间"""
        try:
            if os.path.exists(self.checkpoint_file_path):
                with open(self.checkpoint_file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                return data.get('checkpoint_time', datetime.now().isoformat())
        except:
            pass
        return datetime.now().isoformat()
    
    def perform_checkpoint(self, flushed_pages: List[Dict[str, Any]] = None) -> int:
        """
        执行检查点
        
        Args:
            flushed_pages: 已刷新的页面信息
            
        Returns:
            检查点LSN
        """
        with self.wal_lock:
            # 强制刷新WAL
            self._flush_buffer()
            
            checkpoint_lsn = self.current_lsn - 1
            checkpoint_time = datetime.now()
            
            # 创建检查点信息
            checkpoint_info = {
                'checkpoint_lsn': checkpoint_lsn,
                'checkpoint_time': checkpoint_time.isoformat(),
                'flushed_pages': flushed_pages or [],
                'active_transactions': list(self.active_transactions.keys()),
                'wal_files_before_checkpoint': self.list_wal_files()
            }
            
            try:
                # 保存检查点信息
                with open(self.checkpoint_file_path, 'w', encoding='utf-8') as f:
                    json.dump(checkpoint_info, f, indent=2, default=str)
                
                # 更新控制信息
                self._save_control_info()
                
                self.logger.info("WAL_CHECKPOINT", checkpoint_lsn, f"Checkpoint completed at LSN {checkpoint_lsn}")
                
                return checkpoint_lsn
                
            except Exception as e:
                self.logger.error("WAL_CHECKPOINT", 0, f"Failed to perform checkpoint: {e}")
                raise
    
    def get_last_checkpoint_info(self) -> Dict[str, Any]:
        """获取最后检查点信息"""
        try:
            if os.path.exists(self.checkpoint_file_path):
                with open(self.checkpoint_file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
        except Exception as e:
            self.logger.error("WAL_CHECKPOINT_INFO", 0, f"Failed to get checkpoint info: {e}")
        
        return {
            'checkpoint_lsn': 0,
            'checkpoint_time': datetime.now().isoformat(),
            'flushed_pages': [],
            'active_transactions': [],
            'wal_files_before_checkpoint': []
        }
    
    def close(self):
        """关闭WAL管理器"""
        try:
            # 刷新剩余缓冲区
            self.force_flush()
            
            # 关闭WAL文件
            if self.current_wal_file:
                self.current_wal_file.close()
                self.current_wal_file = None
            
            self.logger.info("WAL_CLOSE", 0, "WAL Manager closed")
            
        except Exception as e:
            self.logger.error("WAL_CLOSE", 0, f"Error closing WAL Manager: {e}")
    
    def __enter__(self):
        """支持with语句"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持with语句"""
        self.close()
